package jwt

import (
	"context"
	"github.com/go-kratos/kratos/v2/errors"
	"github.com/go-kratos/kratos/v2/middleware"
	"github.com/go-kratos/kratos/v2/middleware/auth/jwt"
	"github.com/go-kratos/kratos/v2/transport"
	jwtV4 "github.com/golang-jwt/jwt/v4"
	"google.golang.org/grpc/metadata"
	"strings"
	"time"
)

const (

	// bearerWord the bearer key word for authorization
	bearerWord string = "Bearer"

	// bearerFormat authorization token format
	bearerFormat string = "Bearer %s"

	// authorizationKey holds the key used to store the JWT Token in the request tokenHeader.
	authorizationKey string = "Authorization"

	// reason holds the error reason.
	reason string = "UNAUTHORIZED"
)

var (
	ErrAuthFail = errors.Unauthorized(reason, "JWT auth is fail")
)

type AuthInfo struct {
	UserId int64 `json:"user_id"`
}

type MyClaims struct {
	AuthInfo
	jwtV4.RegisteredClaims
}

func Auth(key string, info AuthInfo, expire time.Time) (string, error) {
	myClaims := MyClaims{
		info,
		jwtV4.RegisteredClaims{
			ExpiresAt: jwtV4.NewNumericDate(expire), //设置JWT过期时间,此处设置为2小时
		},
	}
	claims := jwtV4.NewWithClaims(jwtV4.SigningMethodHS256, myClaims)
	//加盐
	return claims.SignedString([]byte(key))
}

// CheckJWT 解析
func CheckJWT(key string, jwtToken string) (map[string]interface{}, error) {
	token, err := jwtV4.Parse(jwtToken, func(token *jwtV4.Token) (interface{}, error) {
		return []byte(key), nil
	})
	if err != nil {
		return nil, err
	}
	if claims, ok := token.Claims.(jwtV4.MapClaims); ok {
		result := make(map[string]interface{}, 2)
		result["user_id"] = claims["user_id"]
		return result, nil
	} else {
		return nil, jwt.ErrTokenInvalid
	}
}

// NewAuthServer jwt Server中间件
func NewAuthServer(key string) func(handler middleware.Handler) middleware.Handler {
	return func(handler middleware.Handler) middleware.Handler {
		return func(ctx context.Context, req interface{}) (reply interface{}, err error) {
			var jwtToken string
			if md, ok := metadata.FromIncomingContext(ctx); ok {
				jwtToken = md.Get("x-md-global-jwt")[0]
			} else if header, ok := transport.FromServerContext(ctx); ok {
				auths := strings.SplitN(header.RequestHeader().Get(authorizationKey), " ", 2)
				if len(auths) != 2 || !strings.EqualFold(auths[0], bearerWord) {
					return nil, jwt.ErrMissingJwtToken
				}
				jwtToken = auths[1]
			} else {
				// 缺少可认证的token，返回错误
				return nil, ErrAuthFail
			}
			token, err := CheckJWT(key, jwtToken)
			if err != nil {
				// 缺少合法的token，返回错误
				return nil, ErrAuthFail
			}
			ctx = context.WithValue(ctx, "user_id", int64(token["user_id"].(float64)))
			reply, err = handler(ctx, req)
			return
		}
	}
}

// NewAuthClient jwt Client中间件
func NewAuthClient(key string) middleware.Middleware {
	return func(handler middleware.Handler) middleware.Handler {
		return func(ctx context.Context, req interface{}) (interface{}, error) {

			if header, ok := transport.FromServerContext(ctx); ok {
				//如果Token为空则生成token，如果
				if header.RequestHeader().Get(authorizationKey) != "" {
					auths := strings.SplitN(header.RequestHeader().Get(authorizationKey), " ", 2)
					if len(auths) != 2 || !strings.EqualFold(auths[0], bearerWord) {
						return nil, jwt.ErrMissingJwtToken
					}
					jwtToken := auths[1]
					//元数据传递
					ctx = metadata.AppendToOutgoingContext(ctx, "x-md-global-jwt", jwtToken)
				}
			}
			return handler(ctx, req)
		}
	}
}
